import time
import sys
import argparse
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient

def move_robot(vx=0.0, vy=0.0, vyaw=0.0, network_interface="eth0", duration=3.0):
    """
    控制G1机器人移动
    
    参数:
        vx (float): 前进/后退速度，正值前进，负值后退，范围[-1.0, 1.0]
        vy (float): 左右移动速度，正值向左，负值向右，范围[-1.0, 1.0]
        vyaw (float): 旋转速度，正值逆时针，负值顺时针，范围[-1.0, 1.0]
        network_interface (str): 网络接口名称，默认为eth0
        duration (float): 移动持续时间（秒）
    """
    print(f"初始化机器人连接，网络接口: {network_interface}")
    print(f"移动参数: 前进/后退={vx}, 左右={vy}, 旋转={vyaw}, 持续时间={duration}秒")
    
    try:
        # 初始化网络通道
        ChannelFactoryInitialize(0, network_interface)
        
        # 创建运动控制客户端
        sport_client = LocoClient()  
        sport_client.SetTimeout(10.0)
        sport_client.Init()
        
        # 确保机器人处于站立状态
        print("确保机器人站立...")
        sport_client.Squat2StandUp()
        time.sleep(2.0)
        
        # 执行移动命令
        print("开始移动...")
        # 直接使用SetVelocity而不是Move，因为Move函数内部会覆盖duration参数
        sport_client.SetVelocity(vx, vy, vyaw, duration)
        
        # 等待指定时间
        time.sleep(duration)
        
        # 停止移动
        print("停止移动...")
        sport_client.SetVelocity(0.0, 0.0, 0.0, 1.0)
        
        print("移动完成")
    except Exception as e:
        print(f"发生错误: {e}")
        print("提示: 请确保机器人已连接且网络接口正确")
        print("      机器人可能需要先站立才能接收移动命令")

def parse_set_velocity(velocity_str):
    """
    解析速度参数字符串，格式为"vx vy vyaw duration"
    
    参数:
        velocity_str (str): 速度参数字符串，例如"0.5 0 0 1"
        
    返回:
        tuple: (vx, vy, vyaw, duration)
    """
    try:
        parts = velocity_str.split()
        if len(parts) >= 4:
            return float(parts[0]), float(parts[1]), float(parts[2]), float(parts[3])
        elif len(parts) == 3:
            return float(parts[0]), float(parts[1]), float(parts[2]), 3.0  # 默认持续时间3秒
        elif len(parts) == 2:
            return float(parts[0]), float(parts[1]), 0.0, 3.0  # 默认无旋转，持续3秒
        elif len(parts) == 1:
            return float(parts[0]), 0.0, 0.0, 3.0  # 默认只前进，无侧移，无旋转，持续3秒
        else:
            return 0.3, 0.0, 0.0, 3.0  # 默认值
    except ValueError:
        print(f"警告: 无法解析速度参数 '{velocity_str}'，使用默认值")
        return 0.3, 0.0, 0.0, 3.0  # 默认值

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="G1机器人移动控制")
    parser.add_argument("--network_interface", default="eth0", help="网络接口名称，默认为eth0")
    parser.add_argument("--vx", type=float, default=0.3, help="前进/后退速度，正值前进，负值后退，范围[-1.0, 1.0]，默认0.3")
    parser.add_argument("--vy", type=float, default=0.0, help="左右移动速度，正值向左，负值向右，范围[-1.0, 1.0]，默认0.0")
    parser.add_argument("--vyaw", type=float, default=0.0, help="旋转速度，正值逆时针，负值顺时针，范围[-1.0, 1.0]，默认0.0")
    parser.add_argument("--duration", type=float, default=1.0, help="移动持续时间（秒），默认1.0秒")
    parser.add_argument("--continuous", action="store_true", help="持续移动模式，直到手动停止")
    parser.add_argument("--set_velocity", type=str, help="一次性设置所有速度参数，格式为'vx vy vyaw duration'，例如'0.5 0 0 1'")
    parser.add_argument("--no_stand", action="store_true", help="不执行站立动作，直接移动")
    
    args = parser.parse_args()
    
    # 检查网络接口参数是否包含@符号（适配命令行调用方式）
    if "@" in args.network_interface:
        args.network_interface = args.network_interface.split("@")[0]
    
    # 如果提供了set_velocity参数，解析它
    if args.set_velocity:
        vx, vy, vyaw, duration = parse_set_velocity(args.set_velocity)
    else:
        vx, vy, vyaw, duration = args.vx, args.vy, args.vyaw, args.duration
    
    # 如果是持续移动模式，设置较长的持续时间
    if args.continuous:
        duration = 864000.0  # 10天，近似于持续移动
    
    # 如果指定了--no_stand参数，则修改move_robot函数内部逻辑跳过站立步骤
    if args.no_stand:
        # 创建一个临时函数，不执行站立动作
        def move_robot_no_stand(vx=0.0, vy=0.0, vyaw=0.0, network_interface="eth0", duration=3.0):
            print(f"初始化机器人连接，网络接口: {network_interface}")
            print(f"移动参数: 前进/后退={vx}, 左右={vy}, 旋转={vyaw}, 持续时间={duration}秒")
            
            try:
                # 初始化网络通道
                ChannelFactoryInitialize(0, network_interface)
                
                # 创建运动控制客户端
                sport_client = LocoClient()  
                sport_client.SetTimeout(10.0)
                sport_client.Init()
                
                # 执行移动命令
                print("开始移动...")
                sport_client.SetVelocity(vx, vy, vyaw, duration)
                
                # 等待指定时间
                time.sleep(duration)
                
                # 停止移动
                print("停止移动...")
                sport_client.SetVelocity(0.0, 0.0, 0.0, 1.0)
                
                print("移动完成")
            except Exception as e:
                print(f"发生错误: {e}")
                print("提示: 请确保机器人已连接且网络接口正确")
        
        # 使用不执行站立动作的函数
        move_robot_no_stand(vx, vy, vyaw, args.network_interface, duration)
    else:
        # 使用原始函数，包含站立动作
        move_robot(vx, vy, vyaw, args.network_interface, duration)